#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File          :    test.py    
@Contact       :    LJL959@QQ.COM
@License       :    (C)Copyright 2020-2021, Liugroup-NLPR-CASIA
@Modify Time   :    2020/12/9 14:49        
@Author        :    LiuJiaolong
@Version       :    1.0
@Description    :   测试网络
"""
# import lib
import torch
import numpy as np
import torch.nn as nn

from D2L.MyFirstDemoWithFMRIData.network import Net


model_file = './model/model.pth'        # 训练所存储的模型
# train_data      # 训练数据集
data_dir = './data/testData/'
n_epoch = 6
batch_size = 3
# cnt = 0
all_test_data = []
acc = []

def test():
    model = Net()           # 实例化一个网络
    model                   # 有GPU的后面加.cuda()
    # model = nn.DataParallel(model)  # 多GPU并行训练
    model.load_state_dict(torch.load(model_file), False)    # 加载训练好的模型，若出错，可能是训练时使用的环境和测试时的环境不一致导致，加个False即可
    model.eval()            # 执行一个字符串表达式，并返回表达式的值
    print(model)
    cnt = 0
    data_with_label = np.load(data_dir + '/drug_health_data.npy')   # 加载用于测试的数据
    print(data_with_label.shape)
    criterion = torch.nn.CrossEntropyLoss()                         # 定义loss计算方法，交叉熵，可以理解为两者数值越接近其值越小
    eval_loss = 0.
    eval_acc = 0.

    for epoch in range(n_epoch):  # 读取数据集进行训练
        for batch_idx, data in enumerate(data_with_label):  # 从封装好的npy文件中读取数据与对应标签
            label_temp = data_with_label[batch_idx, 116, 0]  # 标签只占一个位置,在第116行 第0列位置处
            label_x = torch.tensor(label_temp, dtype=torch.long)  # 数据类型转换为tensor

            # # 从封装好的文件里取得数据有问题
            # data_temp = data_with_label[batch_idx, 0:116, 0:150]  # 前0~116存放结点数据，0~150存放特征，所以数据就是0~150全取
            # data_x = torch.tensor(data_temp)  # 转换成tensor
            # data_x_unsqueeze = torch.unsqueeze(data_x, 0)  # 扩充维度
            # data_new = torch.unsqueeze(data_x_unsqueeze, 0)  # 扩充维度


            # 老方法取的数据没问题


            single_data = data[:116]
            single_data = torch.Tensor(single_data)
            all_test_data.append(single_data)
            temp_data = torch.unsqueeze(all_test_data[0], 0)  # 在第一维上扩一维
            new_data = torch.unsqueeze(temp_data, 0)

            out = model(new_data)  # 计算网络输出值，就是输入网络的一个特征数据，输出病人和健康人的概率，调用了网络中的forward方法
            loss = criterion(out, label_x.unsqueeze(0))  # 计算损失，就是网络输出值和实际label的差异，差异越小，说明拟合效果越好
            cnt += 1
            acc_temp = 1 - loss / batch_size
            acc.append(acc_temp.item())
            # print('Epoch:{0}, Frame:{1}, test_loss:{2}'.format(epoch, cnt*batch_size, loss/batch_size))    # 打印每个batch_size的训练结果
    mean_acc = np.array(acc).mean()         # 求均值

    print("mean_acc", mean_acc)


if __name__ == '__main__':
    test()
